import requests
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional, Dropout
from tensorflow.keras.callbacks import ReduceLROnPlateau
import tushare as ts
import math
import random

# 设置随机种子，保证结果的可复现性
random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

# 设置中文字体为黑体，解决负号显示问题（用于部分和matplotlib相关的操作，比如后续模型训练结果可视化等）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 获取股票数据的函数，使用Tushare的数据接口（需替换token才能正常获取数据）
def get_stock_data(token, ts_code, start_date, end_date, fields):
    """
    从tushare获取指定股票代码、日期范围以及字段的股票数据

    参数:
    token (str): 在tushare平台注册获取的token
    ts_code (str): 股票代码，如 "01211.HK"
    start_date (str): 开始日期，格式为 "YYYYMMDD"
    end_date (str): 结束日期，格式为 "YYYYMMDD"
    fields (list): 要获取的股票数据字段列表，如 ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    返回:
    pd.DataFrame: 获取到的股票数据，如果获取失败返回None
    """
    try:
        pro = ts.pro_api(token)  # 请替换为自己的Tushare token
        data = pro.daily(ts_code=ts_code,
                         start_date=start_date,
                         end_date=end_date,
                         fields=fields)
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None

# 数据预处理函数，进行缺失值处理、数据标准化等操作，增加更多详细信息打印
def preprocess_data(data):
    df = data.copy()
    # 将交易日期列设置为索引，并转换为日期时间类型
    df.set_index('trade_date', inplace=True)
    df.index = pd.to_datetime(df.index)
    print("设置日期索引后的基本信息：")
    print(df.info())

    # 去除包含缺失值的行
    df.dropna(inplace=True)
    print("去除缺失值后的基本信息：")
    print(df.info())

    # 提取需要的列（开盘价、最高价、最低价、收盘价、成交量）
    columns = ['open', 'high', 'low', 'close', 'vol', 'pre_close', 'change', 'pct_chg']
    df = df[columns]
    print("提取关键列后的基本信息：")
    print(df.info())

    # 数据标准化，这里使用StandardScaler将数据标准化
    scaler = StandardScaler()
    if df.shape[0] == 0:
        print("数据为空，无法进行标准化处理！")
        exit(1)

    df_scaled = scaler.fit_transform(df)
    print("标准化后的数据形状:", df_scaled.shape)
    return df, df_scaled, scaler

# 创建训练数据集和测试数据集的函数，按照时间序列划分数据
def create_dataset(dataset, time_steps=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - time_steps - 1):
        a = dataset[i:(i + time_steps), :]
        dataX.append(a)
        dataY.append(dataset[i + time_steps, -1])
    return np.array(dataX), np.array(dataY)

# 构建并训练LSTM模型的函数，支持不同模型类型选择
def build_and_train_model(model_type, X_train, y_train, X_test, y_test):
    model = Sequential()
    if model_type == 1:
        # 单层 LSTM
        model.add(LSTM(units=100, activation='relu',
                       input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(units=1))
    elif model_type == 2:
        # 多层 LSTM
        model.add(LSTM(units=100, activation='relu', return_sequences=True,
                       input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dropout(0.2))  # 加入Dropout层
        model.add(LSTM(units=100, activation='relu'))
        model.add(Dense(1))
    elif model_type == 3:
        # 双向 LSTM
        model.add(Bidirectional(LSTM(100, activation='relu'),
                                input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(1))

    model.compile(loss='mean_squared_error', optimizer='adam')

    # 学习率动态调整
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, min_lr=1e-6)
    history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=100, batch_size=32, callbacks=[reduce_lr], verbose=1)

    return model, history

# 计算均方根误差（RMSE），用于评估模型性能
def calculate_rmse(y_true, y_pred):
    min_len = min(len(y_true), len(y_pred))  # 获取较小的长度
    y_true = y_true[:min_len]  # 修剪 y_true
    y_pred = y_pred[:min_len]  # 修剪 y_pred

    return math.sqrt(((y_true - y_pred) ** 2).mean())

if __name__ == "__main__":
    # Tushare平台获取的token，需替换为你自己的真实token
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'
    ts_code = "600519.SH"  # 贵州茅台股票代码，可替换为其他感兴趣的股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)

    # 检查获取的数据是否为空
    if stock_data is None or stock_data.empty:
        print("没有获取到有效的股票数据！")
        exit(1)

    # 数据预处理
    df, df_scaled, scaler = preprocess_data(stock_data)

    # 划分训练集和测试集
    print("数据总长度:", len(df_scaled))
    train_size = int(len(df_scaled) * 0.8)
    print("计算出的训练集预期大小:", train_size)
    test_size = len(df_scaled) - train_size
    print("计算出的测试集预期大小:", test_size)
    train_data, test_data = df_scaled[0:train_size, :], df_scaled[train_size:len(df_scaled), :]
    print("实际划分后训练集数据形状:", train_data.shape)
    print("实际划分后测试集数据形状:", test_data.shape)
    time_steps = 20  # 增加时间步长
    X_train, y_train = create_dataset(train_data, time_steps)
    X_test, y_test = create_dataset(test_data, time_steps)
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], df.shape[1]))
    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], df.shape[1]))

    # 选择模型类型（1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM）
    model_type = 2  # 尝试使用多层 LSTM
    model, history = build_and_train_model(model_type, X_train, y_train, X_test, y_test)

    # 使用训练好的模型进行预测
    train_predict = model.predict(X_train)
    train_predict = np.reshape(train_predict, (train_predict.shape[0], 1))

    # 对训练集进行反标准化处理
    train_data_trimmed = train_data[time_steps:len(train_predict) + time_steps, 0:df.shape[1] - 1]
    train_predict = scaler.inverse_transform(np.concatenate((train_data_trimmed, train_predict), axis=1))[:, -1]

    # 使用模型进行预测并计算RMSE
    test_predict = model.predict(X_test)
    test_predict = np.reshape(test_predict, (test_predict.shape[0], 1))

    test_data_trimmed = test_data[time_steps:len(test_predict) + time_steps, 0:df.shape[1] - 1]
    test_predict = scaler.inverse_transform(np.concatenate((test_data_trimmed, test_predict), axis=1))[:, -1]

    rmse = calculate_rmse(y_test, test_predict)
    print(f"测试集RMSE: {rmse}")

    # 使用Plotly绘制真实股价与预测股价对比
    fig = go.Figure()

    fig.add_trace(go.Scatter(x=df.index[train_size + time_steps + 1:], y=y_test, mode='lines', name='真实值'))
    fig.add_trace(go.Scatter(x=df.index[train_size + time_steps + 1:], y=test_predict, mode='lines', name='预测值'))

    fig.update_layout(title="真实值与预测值对比",
                      xaxis_title="日期",
                      yaxis_title="收盘价")
    fig.show(renderer="browser")  # 强制在浏览器中显示图表
